#include <omplapp/config.h>
#include <ompl/tools/benchmark/Benchmark.h>
#include <omplapp/geometry/GeometrySpecification.h>
#include "RRTIterationCount.h"
#include "RRTConnectIterationCount.h"
#include "LazyRRTIterationCount.h"
#include "pRRTIterationCount.h"
#include "FCLStateValidityCheckerCounter.h"
#include <omplapp/apps/SE3RigidBodyPlanning.h>
#include <ompl/base/samplers/UniformValidStateSampler.h>
#include <boost/lexical_cast.hpp>
#include <stdio.h>
#include "PoissonRRT.h"
#include "PoissonSampler.h"
#include "CUDAPoissonRRT.h"
#include "CUDARRTConnect.h"

using namespace ompl;

base::ValidStateSamplerPtr allocUniformStateSampler3D(const base::SpaceInformation *si)
{
	return base::ValidStateSamplerPtr(new base::UniformValidStateSampler(si));
}

void benchmark3d0(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("easy");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Easy_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Easy_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	base::ScopedState<base::SE3StateSpace> goal(start);

	// random polygons
	start->setX(270);
	start->setY(160);
	start->setZ(-200);
	start->rotation().setIdentity();
	goal->setX(270);
	goal->setY(160);
	goal->setZ(-400);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 457.960449219);
	bounds.setHigh(1, 321.25);
	bounds.setHigh(2, -72.8550872803);
	bounds.setLow(0, 14.4604492188);
	bounds.setLow(1, -24.25);
	bounds.setLow(2, -504.855102539);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d1(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("cubicles");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/cubicles_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/cubicles_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	start->setX(-4.96);
	start->setY(-40.62);
	start->setZ(70.57);
	start->rotation().setIdentity();

	base::ScopedState<base::SE3StateSpace> goal(start);
	goal->setX(200.49);
	goal->setY(-40.62);
	goal->setZ(70.57);
	goal->rotation().setIdentity();

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d2(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("alpha-1.5");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/alpha_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/alpha_env-1.5.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	base::ScopedState<base::SE3StateSpace> goal(start);

	// random polygons
	start->setX(-21.91);
	start->setY(-4.11);
	start->setZ(-14.14);
	start->rotation().setIdentity();
	goal->setX(-21.91);
	goal->setY(-4.11);
	goal->setZ(68.86);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 189.05);
	bounds.setHigh(1, 189.18);
	bounds.setHigh(2, 174.86);
	bounds.setLow(0, -281.64);
	bounds.setLow(1, -119.64);
	bounds.setLow(2, -176.86);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d3(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("apartment_piano");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Apartment_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Apartment_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	base::ScopedState<base::SE3StateSpace> goal(start);

	// random polygons
	start->setX(241.81);
	start->setY(106.15);
	start->setZ(36.46);
	start->rotation().setAxisAngle(0, 0, -1, 3.12413936107);
	goal->setX(-31.19);
	goal->setY(-99.85);
	goal->setZ(36.46);
	goal->rotation().setAxisAngle(0, 0, -1, 3.12413936107);

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 295.77);
	bounds.setHigh(1, 168.26);
	bounds.setHigh(2, 90.39);
	bounds.setLow(0, -73.76);
	bounds.setLow(1, -179.59);
	bounds.setLow(2, -0.03);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d4(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("Home");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Home_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Home_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	base::ScopedState<base::SE3StateSpace> goal(start);

	// random polygons
	start->setX(252.95);
	start->setY(-214.95);
	start->setZ(46.19);
	start->rotation().setIdentity();
	goal->setX(262.95);
	goal->setY(75.05);
	goal->setZ(46.19);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 324.997131348);
	bounds.setHigh(1, 337.893371582);
	bounds.setHigh(2, 142.332290649);
	bounds.setLow(0, -383.802642822);
	bounds.setLow(1, -371.469055176);
	bounds.setLow(2, -0.196851730347);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d5(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("Twistycool");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Twistycool_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/Twistycool_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	start->setX(270.);
	start->setY(160.);
	start->setZ(-200.);
	start->rotation().setIdentity();

	base::ScopedState<base::SE3StateSpace> goal(start);
	goal->setX(270.);
	goal->setY(160.);
	goal->setZ(-400.);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 350.);
	bounds.setHigh(1, 250.);
	bounds.setHigh(2, -150.);
	bounds.setLow(0, 200.);
	bounds.setLow(1, 75.);
	bounds.setLow(2, -450.);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d6(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("bugtrap");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/bugtrap_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/bugtrap_env.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	start->setX(17.18);
	start->setY(0.89);
	start->setZ(-4.62);
	start->rotation().setIdentity();

	base::ScopedState<base::SE3StateSpace> goal(start);
	goal->setX(45.18);
	goal->setY(0.89);
	goal->setZ(-4.62);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 56.99);
	bounds.setHigh(1, 19.57);
	bounds.setHigh(2, 15.27);
	bounds.setLow(0, 4.04);
	bounds.setLow(1, -44.49);
	bounds.setLow(2, -13.27);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void benchmark3d7(std::string& benchmark_name, app::SE3RigidBodyPlanning& setup,
	double& runtime_limit, double& memory_limit, int& run_count)
{
	benchmark_name = std::string("alpah-1.2");
	std::string robot_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/alpha_robot.dae";
	std::string env_fname = std::string(OMPLAPP_RESOURCE_DIR) + "/3D/alpha_env-1.2.dae";
	setup.setRobotMesh(robot_fname.c_str());
	setup.setEnvironmentMesh(env_fname.c_str());

	// define starting & goal state
	base::ScopedState<base::SE3StateSpace> start(setup.getSpaceInformation());
	start->setX(-21.91);
	start->setY(-11.11);
	start->setZ(-14.14);
	start->rotation().setIdentity();

	base::ScopedState<base::SE3StateSpace> goal(start);
	goal->setX(-21.91);
	goal->setY(-11.11);
	goal->setZ(48.86);
	goal->rotation().setIdentity();

	base::RealVectorBounds bounds(3);
	bounds.setHigh(0, 189.05);
	bounds.setHigh(1, 189.18);
	bounds.setHigh(2, 174.86);
	bounds.setLow(0, -281.64);
	bounds.setLow(1, -119.64);
	bounds.setLow(2, -176.86);
	setup.getStateSpace()->as<base::SE3StateSpace>()->setBounds(bounds);

	setup.setStartAndGoalStates(start, goal);
	setup.getSpaceInformation()->setStateValidityCheckingResolution(0.01);
	setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
	setup.setup();

	std::vector<double> cs(3);
	cs[0] = 35; cs[1] = 35; cs[2] = 35;
	setup.getStateSpace()->getDefaultProjection()->setCellSizes(cs);
}

void preRunEvent3D(const base::PlannerPtr &planner)
{
}

void postRunEvent3D(const base::PlannerPtr &planner, tools::Benchmark::RunProperties &run)
{

}

int main(int argc, char **argv)
{
	ompl::RNG::setSeed(1);

	app::SE3RigidBodyPlanning setup;
	std::string benchmark_name;
	double runtime_limit, memory_limit;
	int run_count;

	int benchmark_id = 1;
	int sampler_id = 2;
	int num_threads = 1;

	bool parallelOnly = false;

	if (argc > 1)
	{
		benchmark_id = boost::lexical_cast<int>(argv[1]);
	}
	if (argc > 2)
	{
		sampler_id = boost::lexical_cast<int>(argv[2]);
	}
	if (argc > 3)
	{
		num_threads = boost::lexical_cast<int>(argv[3]);
	}
	if (argc > 4)
	{
		parallelOnly = true;
	}

	switch (benchmark_id)
	{
	case 1:
		benchmark3d1(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 2:
		benchmark3d2(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;
	case 3:
		benchmark3d3(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 4:
		benchmark3d4(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 5:
		benchmark3d5(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 6:
		benchmark3d6(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 7:
		benchmark3d7(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;

	case 0:
	default:
		benchmark3d0(benchmark_name, setup, runtime_limit, memory_limit, run_count);
		break;
	}

	runtime_limit = 120.0;
	memory_limit = 10000.0; // set high because memory usage is not always estimated correctly
	run_count = 1;


	// create the benchmark object and add all the planners we'd like to run
	tools::Benchmark::Request request(runtime_limit, memory_limit, run_count);
	tools::Benchmark b(setup, benchmark_name);

	request.useThreads = false;

	// optionally set pre & pos run events
	b.setPreRunEvent(boost::bind(&preRunEvent3D, _1));
	b.setPostRunEvent(boost::bind(&postRunEvent3D, _1, _2));

	if (!parallelOnly)
	{
		switch (sampler_id)
		{
		case 3:
			b.addPlanner(base::PlannerPtr(new CUDARRTConnect(setup.getSpaceInformation(), &setup.getGeometrySpecification())));
			setup.getSpaceInformation()->setValidStateSamplerAllocator(poisson::PoissonSampler::allocPoissonStateSampler);
			b.setExperimentName(benchmark_name + "_poisson_sampler");
			break;

		case 2:
			b.addPlanner(base::PlannerPtr(new poisson::CUDAPoissonRRT(setup.getSpaceInformation(), &setup.getGeometrySpecification())));
			setup.getSpaceInformation()->setValidStateSamplerAllocator(poisson::PoissonSampler::allocPoissonStateSampler);
			b.setExperimentName(benchmark_name + "_poisson_sampler");
			break;

		case 1:
			b.addPlanner(base::PlannerPtr(new poisson::PoissonRRT(setup.getSpaceInformation())));
			setup.getSpaceInformation()->setValidStateSamplerAllocator(poisson::PoissonSampler::allocPoissonStateSampler);
			b.setExperimentName(benchmark_name + "_poisson_sampler");
			break;

		default:
			b.addPlanner(base::PlannerPtr(new RRTIterationCount(setup.getSpaceInformation())));
			b.addPlanner(base::PlannerPtr(new RRTConnectIterationCount(setup.getSpaceInformation())));
			b.addPlanner(base::PlannerPtr(new LazyRRTIterationCount(setup.getSpaceInformation())));
			setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
			b.setExperimentName(benchmark_name + "_uniform_sampler");
			break;
		}
	}
	else
	{
		const int THREAD_COUNT[] = { 1, 2, 4, 8, 16, 24, 32 };
		if (sampler_id == 1)
		{

		}
		else
		{
			for (int i = 0; i < 7; ++i)
			{
				b.addPlanner(base::PlannerPtr(new pRRTIterationCount(setup.getSpaceInformation(), THREAD_COUNT[i])));
			}
			setup.getSpaceInformation()->setValidStateSamplerAllocator(&allocUniformStateSampler3D);
		}
		b.setExperimentName(benchmark_name + "_parallel");
	}

	b.benchmark(request);

	char fileName[80];
#ifdef WIN32
	sprintf_s(fileName, 80, "3dbenchmark%d%d_%d%s.log", benchmark_id, sampler_id, num_threads, parallelOnly ? "p" : "");
#else
	snprintf(fileName, 80, "3dbenchmark%d%d_%d%s.log", benchmark_id, sampler_id, num_threads, parallelOnly ? "p" : "");
#endif
	b.saveResultsToFile(fileName);

	return 0;
}

